/**
* @brief: a copy of file generata_yolt_train_data.hpp for rewrite
* @author: mysticalwing
* @modify: name            time            brief
*          mysticalwing    2018.11.11      create 
*          mysticalwing    2018.11.20      修改产生训练数据的策略，使得有飞行器的图和没有飞行器的图比例为 1:2  
*          mysticalwing    2018.11.22      将数据归一化到[0, 1] 内  
*/

#ifndef GENERATE_YOLT_TRAIN_DATA
#define GENERATE_YOLT_TRAIN_DATA

#define GENERATE_YOLT_TRAIN_DATA_TEST true

#include "../nn_yolt/yolt_def.hpp"
#include "files_io.h"
#include "type_def.h"
#include "visualize_train_data.h"

#include <map>
#include <vector>
#include <string>
#include <iostream>
#include <algorithm>
#include <opencv2/opencv.hpp>
#include <boost/filesystem.hpp>

namespace tools
{

struct Shape3d
{
    int w, h, d;
    Shape3d(const int width, const int height, const int depth = 1)
        : w(width), h(height), d(depth)
    {
        ;
    }

    Shape3d(const i_vec_t &size)
    {
        assert(size.size() == 3 || size.size() == 2);
        w = size[0];
        h = size[1];
        d = (size.size() == 3) ? size[2] : 1;
    }
};

struct BoundingBox
{
    int ox, oy, w, h, cx, cy;

    BoundingBox()
    {
        ox = oy = cx = cy = w = h = 0;
    }

    BoundingBox(const i_vec_t &vertexes)
    {
        assert(vertexes.size() == 8);

        ox = w = vertexes[0];
        oy = h = vertexes[1];
        cx = cy = 0;
        for (int i = 0; i < 8; i += 2)
        {
            ox = std::min(ox, vertexes[i]);
            oy = std::min(oy, vertexes[i + 1]);
            w = std::max(w, vertexes[i]);
            h = std::max(h, vertexes[i + 1]);
        }

        w -= ox;
        h -= oy;
        cx = ox + w / 2 + 0.5;
        cy = oy + h / 2 + 0.5;
    }

    BoundingBox(const struct BoundingBox &right)
        : ox(right.ox), oy(right.oy),
          cx(right.cx), cy(right.cy),
          w(right.w), h(right.h)
    {
        ;
    }

    BoundingBox(const cv::Rect &rect)
        : ox(rect.x), oy(rect.y),
          w(rect.width), h(rect.height)
    {
        cx = rect.x + rect.width / 2 + 0.5;
        cy = rect.y + rect.height / 2 + 0.5;
    }

    bool operator<(const struct BoundingBox &right) const
    {
        return (cy < right.cy) || ((cy == right.cy) && (cx < right.cx));
    }

    bool is_contain_point(const int x, const int y) const
    {
        return (ox <= x) && (x < (ox + w)) &&
               (oy <= y) && (y < (oy + h));
    }

    str_t str() const
    {
        return std::to_string(ox) + "_" + std::to_string(oy) + "_" +
               std::to_string(w) + "_" + std::to_string(h);
    }

    struct BoundingBox &operator=(const struct BoundingBox &right)
    {
        ox = right.ox;
        oy = right.oy;

        cx = right.cx;
        cy = right.cy;

        w = right.w;
        h = right.h;

        return *this;
    }

    friend std::ostream &operator<<(std::ostream &os, const struct BoundingBox &right);
};

std::ostream &operator<<(std::ostream &os, const struct BoundingBox &right)
{
    os << "origin: [" << right.ox << " " << right.oy << "]\n"
       << "center: [" << right.cx << " " << right.cy << "]\n"
       << "size  : [" << right.w << " " << right.h << "]\n";
    return os;
}

class GenerateYoltTrainData
{

    typedef std::map<struct BoundingBox, cv::Mat> image_map_t;

  public:
    GenerateYoltTrainData(const Shape3d &yolt_input_shape,
                          const int _cell_num_hor,
                          const int _cell_num_ver,
                          const int _box_num_per_cell,
                          const float_t _overlap_percentage)
        : image_width(yolt_input_shape.w), image_height(yolt_input_shape.h),
          image_depth(yolt_input_shape.d), box_num_per_cell(_box_num_per_cell),
          cell_num_hor(_cell_num_hor), cell_num_ver(_cell_num_ver),
          overlap_percentage(_overlap_percentage)
    {
        assert((image_width % cell_num_hor) == 0);
        assert((image_height % cell_num_ver) == 0);
        assert(0.0f <= overlap_percentage <= 1.0f);

        cell_width = image_width / cell_num_hor;
        cell_height = image_height / cell_num_ver;
        overlap_hor = overlap_percentage * image_width;
        overlap_ver = overlap_percentage * image_height;
    }

    ~GenerateYoltTrainData()
    {
        ;
    }

    void generate(const str_t &root_dir, const int max_num = -1);
    void handle_single_image(const str_t &large_image_path,
                             const str_t &label_path,
                             const str_t &save_dir);

  private:
    /**
     * @brief: 读取 'big_image_path' 中的图像，并将其裁剪为指定大小，存放在 'small_image_cutouts' 容器中
     */
    void cutting_image(const str_t big_image_path, image_map_t &small_image_cutouts) const;

    void cutting_image_by_vertical(const cv::Mat &image,
                                   image_map_t &image_cutouts,
                                   cv::Rect &roi, const int step_vertical) const;

    f_vec_t calc_label_params(const struct BoundingBox &box_outer, const struct BoundingBox &box_inner) const;

  private:
    const int image_width, image_height, image_depth;
    const float_t overlap_percentage;
    int cell_num_ver, cell_num_hor, box_num_per_cell;
    int cell_width, cell_height;
    int overlap_ver, overlap_hor;
};

void GenerateYoltTrainData::generate(const str_t &root_dir, const int max_num)
{
    // 处理文件路径
    s_vec_t directorise;
    tools::get_dir_directorise(root_dir, directorise);
    str_t images_dir, labels_dir, saved_dir;
    for (int i = 0; i < directorise.size(); i++)
    {
        if (directorise[i].find("my_images") != std::string::npos)
        {
            images_dir = directorise[i];
        }
        else if (directorise[i].find("my_labels") != std::string::npos)
        {
            labels_dir = directorise[i];
        }
        else if (directorise[i].find("yolt") != std::string::npos)
        {
            saved_dir = directorise[i] + "/";
        }
    }

    std::cout << "Images directory: " << images_dir << std::endl;
    std::cout << "Labels directory: " << labels_dir << std::endl;
    std::cout << "Saved directory: " << saved_dir << std::endl;

    // 获得所有图片和标签的路径
    s_vec_t image_paths, label_txt_paths;
    get_dir_file_paths(images_dir, image_paths);
    get_dir_file_paths(labels_dir, label_txt_paths);
    std::sort(image_paths.begin(), image_paths.end());
    std::sort(label_txt_paths.begin(), label_txt_paths.end());

    int n = image_paths.size() > max_num ? max_num : image_paths.size();
    for (int i = 0; i < n; i++)
    {
        std::cout << "Current file name: " << image_paths[i] << std::endl;
        handle_single_image(image_paths[i], label_txt_paths[i], saved_dir);
    }
}

void GenerateYoltTrainData::handle_single_image(const str_t &large_image_path,
                                                const str_t &label_path,
                                                const str_t &save_dir)
{
    const str_t target_name = "plane";

    // 提取目标的坐标和类别信息
    s_vec_t labels;
    boxes_indexes_vec_t boxes_indexes;
    extract_info_from_txt_file(label_path, boxes_indexes, labels);

    // 提取当前图像中所有类别为 'target_name' 的目标的位置信息
    std::vector<struct BoundingBox> target_boxes;
    for (int i = 0; i < labels.size(); i++)
    {
        if (labels[i] == target_name)
        {
            target_boxes.push_back(BoundingBox(boxes_indexes[i]));
        }
    }

    // 当前图片不包含指定类别的目标，则不进行任何处理
    if (target_boxes.empty())
    {
        return;
    }

    // 将大图裁剪为指定大小的小图
    image_map_t small_images;
    cutting_image(large_image_path, small_images);

    // 处理所有小图
    const str_t image_name = tools::get_file_name_from_path(large_image_path);
    const str_t image_format = "." + tools::get_file_format_from_path(large_image_path);
    const str_t txt_format = ".txt";
    nn_yolt::YoltLabel curr_image_label(image_width, image_height, cell_num_hor, cell_num_ver, box_num_per_cell);
    int *box_num_arr = new int[cell_num_hor * cell_num_ver];
    std::memset(box_num_arr, 0, sizeof(int) * cell_num_hor * cell_num_ver);
    bool have_plane = false;
    std::vector<struct BoundingBox> positive_images, negative_images;
    std::vector<nn_yolt::YoltLabel> positive_labels;

    for (auto iter = small_images.begin(); iter != small_images.end(); iter++)
    {
        const struct BoundingBox &curr_image_box = iter->first;
        const cv::Mat &curr_image = iter->second;
        curr_image_label.clear();
        std::memset(box_num_arr, 0, sizeof(int) * cell_num_hor * cell_num_ver);
        have_plane = false;
        // 确定当前图片是否存在 'plane'，并保存相关数据
        for (int i = 0; i < target_boxes.size(); i++)
        {
            const struct BoundingBox &curr_target_box = target_boxes[i];
            if (curr_image_box.is_contain_point(curr_target_box.cx, curr_target_box.cy))
            {
                have_plane = true;
                int cell_nx = (curr_target_box.cx - curr_image_box.ox) / cell_width;
                int cell_ny = (curr_target_box.cy - curr_image_box.oy) / cell_height;
                int box_num_index = cell_ny * cell_num_hor + cell_nx;
                assert (box_num_arr[box_num_index] <= (box_num_per_cell - 1));
                curr_image_label.set(cell_nx, cell_ny, box_num_arr[box_num_index]++,
                                     calc_label_params(curr_image_box, curr_target_box));
            }
        }

        // 归一化数据
        curr_image_label.normalize();

        if (have_plane)
        {
            positive_images.push_back(curr_image_box);
            positive_labels.push_back(curr_image_label);
        }
        else
        {
            negative_images.push_back(curr_image_box);
        }
        // str_t curr_image_name = save_dir + image_name + "_" + curr_image_box.str() + image_format;
        // str_t curr_txt_name = save_dir + image_name + "_" + curr_image_box.str() + txt_format;
        // curr_image_label.save(curr_txt_name);
        // cv::imwrite(curr_image_name, curr_image);
    }

    auto make_file_name = [&](const str_t &name, const str_t &format) -> str_t {
        return save_dir + image_name + "_" + name + format;
    };

    // 保存含有 target 的图片以及其训练数据
    for (int i = 0; i < positive_images.size(); i++)
    {
        positive_labels[i].save(make_file_name(positive_images[i].str(), txt_format));
        cv::imwrite(make_file_name(positive_images[i].str(), image_format), small_images[positive_images[i]]);
    }

    // // 随机选取等数量的不包含 target 的图片作为负样本
    // std::random_shuffle(negative_images.begin(), negative_images.end());
    // int n = positive_images.size() * 0.2;
    // n = n > negative_images.size() ? negative_images.size() : n;

    // curr_image_label.clear();
    // for (int i = 0; i < n; i++)
    // {
    //     curr_image_label.save(make_file_name(negative_images[i].str(), txt_format));
    //     cv::imwrite(make_file_name(negative_images[i].str(), image_format), small_images[negative_images[i]]);
    // }

    delete[] box_num_arr;

    // std::cout << positive_images.size() << std::endl;
    // std::cout << n << std::endl;
}

void GenerateYoltTrainData::cutting_image(const str_t big_image_path, image_map_t &small_image_cutouts) const
{
    int image_read_flag = image_depth == 3 ? cv::IMREAD_COLOR : cv::IMREAD_GRAYSCALE;
    cv::Mat image = cv::imread(big_image_path, image_read_flag);

    assert(image.data);

    const int step_horizontal = image_width - overlap_hor;
    const int step_vertical = image_height - overlap_ver;

    cv::Rect roi(0, 0, image_width, image_height);
    for (roi.x = 0; roi.x < (image.size().width - roi.width); roi.x += step_horizontal)
    {
        cutting_image_by_vertical(image, small_image_cutouts, roi, step_vertical);
    }

    // 处理右边的剩余像素
    roi.x = image.size().width - roi.width;
    cutting_image_by_vertical(image, small_image_cutouts, roi, step_vertical);
}

void GenerateYoltTrainData::cutting_image_by_vertical(const cv::Mat &image,
                                                      image_map_t &image_cutouts,
                                                      cv::Rect &roi, const int step_vertical) const
{
    for (roi.y = 0; roi.y < (image.size().height - roi.height); roi.y += step_vertical)
    {
        image_cutouts.insert(std::pair<BoundingBox, cv::Mat>(roi, cv::Mat(image, roi)));
    }
    roi.y = image.size().height - roi.height;
    image_cutouts.insert(std::pair<BoundingBox, cv::Mat>(roi, cv::Mat(image, roi)));
}

f_vec_t GenerateYoltTrainData::calc_label_params(const struct BoundingBox &box_outer,
                                                 const struct BoundingBox &box_inner) const
{
    f_vec_t res(PARAM_NUM_PER_BOX);

    res[0] = float_t((box_inner.cx - box_outer.ox) % cell_width);
    res[1] = float_t((box_inner.cy - box_outer.oy) % cell_height);
    res[2] = float_t(box_inner.w);
    res[3] = float_t(box_inner.h);
    res[4] = float_t(1.0f);

    return res;
}

////////////////////////////////////////////////////////////////////////////////////////////////////
#if GENERATE_YOLT_TRAIN_DATA_TEST

class TestGenerateYoltTrainData
{
  public:
    static void run()
    {
        const int image_width = 416,
                  image_height = 416,
                  cell_num_hor = 13,
                  cell_num_ver = 13,
                  box_num_per_cell = 4;
        const float_t overlap = 0.15f;

        const str_t save_dir = "../../datas/train/yolt/";

        // 验证数据正确性
        s_vec_t file_paths, label_paths, image_paths;
        tools::get_dir_file_paths(save_dir, file_paths);
        for (int i = 0; i < file_paths.size(); i++)
        {
            str_t file_format = tools::get_file_format_from_path(file_paths[i]);
            if (file_format == "txt")
            {
                label_paths.push_back(file_paths[i]);
            }
            else if (file_format == "png")
            {
                image_paths.push_back(file_paths[i]);
            }
        }

        std::sort(label_paths.begin(), label_paths.end());
        std::sort(image_paths.begin(), image_paths.end());

        cv::Mat curr_image;
        nn_yolt::YoltLabel curr_label(image_width, image_height, cell_num_hor, cell_num_ver, box_num_per_cell);
        f_vec_t params;
        const int cell_width = image_width / cell_num_hor;
        const int cell_height = image_height / cell_num_ver;
        const cv::Scalar color = cv::Scalar(0, 0, 255);
        for (int i = 0; i < image_paths.size(); i++)
        {
            curr_image = cv::imread(image_paths[i]);
            curr_label.load(label_paths[i]);
            curr_label.denormalize();
            std::cout << image_paths[i] << std::endl;
            std::cout << label_paths[i] << std::endl;

            for (int cell_ny = 0; cell_ny < cell_num_ver; cell_ny++)
            {
                for (int cell_nx = 0; cell_nx < cell_num_hor; cell_nx++)
                {
                    for (int box = 0; box < box_num_per_cell; box++)
                    {
                        params = curr_label.get(cell_nx, cell_ny, box);
                        if (int(params[4]) != 0) // iou != 0
                        {
                            int cx = cell_width * cell_nx + params[0];
                            int cy = cell_height * cell_ny + params[1];
                            cv::rectangle(curr_image, cv::Rect(cx - params[2] / 2, cy - params[3] / 2, params[2], params[3]), color, 2);
                        }
                    }
                }
            }

            cv::imshow("Image", curr_image);
            cv::waitKey(0);
        }
    }
};

#endif

} // namespace tools
#endif // !GENERATE_YOLT_TRAIN_DATA